# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： gen_gap_mask
Description :
Author : 'li'
date： 2021/6/30
-------------------------------------------------
Change Activity:
2021/6/30:
-------------------------------------------------
"""
import numpy as np
import cv2
from ml.cv.algorithms.sort_points.sort_points import resort_points
from ml.cv.util.distance import get_distance
from ml.cv.util.line_util import get_k_b


class GapMaskFormat:
    def __init__(self, gap_mask, threshold=0.6, intersection_threshold=1.4):
        """

        Args:
            gap_mask:
            threshold:
            intersection_threshold:
        """

        self.original_gap_mask = gap_mask
        self.threshold = threshold
        self.intersection_threshold = intersection_threshold

    @staticmethod
    def _get_connected_domain(bin_mask, connectivity=8):
        """
        get connected domain
        """
        bin_mask = bin_mask.astype(np.uint8) * 255
        _, labels = cv2.connectedComponents(bin_mask, connectivity=connectivity)
        return labels

    @staticmethod
    def _get_max_size_contours(contours):
        """
        get max size contours
        :param contours:
        :return:
        """
        return_contours = contours[0]
        length = len(contours)
        for i in range(1, length):
            ret_size = return_contours.shape[0]
            new_size = contours[i].shape[0]
            if ret_size < new_size:
                return_contours = contours[i]
        return return_contours

    @staticmethod
    def _get_line_by_bbox(box):
        box = resort_points(box)
        if get_distance(box[0], box[1]) > get_distance(box[1], box[2]):
            line = [(box[0] + box[3]) / 2, (box[2] + box[1]) / 2]
        else:
            line = [(box[0] + box[3]) / 2, (box[2] + box[1]) / 2]
        return line

    @staticmethod
    def _draw_line_mask(lines, format_mask):
        line_mask = np.zeros_like(format_mask)
        for line in lines:
            p0 = tuple(line[0].astype(np.int32).tolist())
            p1 = tuple(line[1].astype(np.int32).tolist())
            cv2.line(line_mask, p0, p1, 1, 1)
            line[0][1] = line[0][1] + 1
            line[1][1] = line[1][1] + 1
            p0 = tuple(line[0].astype(np.int32).tolist())
            p1 = tuple(line[1].astype(np.int32).tolist())
            cv2.line(line_mask, p0, p1, 1, 1)
        return line_mask

        pass

    def _load_gap_line_mask(self):
        format_mask = np.zeros_like(self.original_gap_mask)
        format_mask[np.where(self.original_gap_mask > self.threshold)] = 1
        connected_domain = self._get_connected_domain(format_mask)
        max_domain = connected_domain.max()
        if max_domain == 0:
            return None
        lines = []
        for i in range(1, max_domain + 1):
            location = np.where(connected_domain == i)
            points_size = location[0].size
            if points_size < 5:
                continue
            bg = np.zeros_like(connected_domain, dtype=np.uint8)
            bg[location] = 255
            contours, _ = cv2.findContours(bg, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
            contours = self._get_max_size_contours(contours)
            rect = cv2.minAreaRect(contours)
            box = cv2.boxPoints(rect)
            line = self._get_line_by_bbox(box)
            line = self.extend_line(line)
            line = line + 1
            line = self.extend_line(line)
            lines.append(line)
        if len(lines) == 0:
            return None
        line_mask = self._draw_line_mask(lines, format_mask)
        return line_mask

    @staticmethod
    def extend_line(line, pixel_size=5):
        """

        """
        line = np.array(line)
        length = get_distance(line[0], line[1])
        if length < 4:
            return line
        k, b = get_k_b(line[0], line[1])
        if abs(k) > 1:
            return line
        if line[0][0] < line[1][0]:
            p0 = line[0]
            p1 = line[1]
        else:
            p1 = line[0]
            p0 = line[1]
        new_p0 = [p0[0] - pixel_size, k * (p0[0] - pixel_size) + b]
        new_p1 = [p1[0] + pixel_size, k * (p1[0] + pixel_size) + b]
        return np.array([new_p0, new_p1])

    def format_gap_mask(self):
        return self._load_gap_line_mask()
